import mss
import mss.tools
import tkinter as tk
from tkinter import simpledialog


def full_screen_screenshot(output="full_screen_screenshot.png"):
    """
        全屏截图
    """
    """Capture a screenshot of the entire screen."""
    with mss.mss() as sct:
        screenshot = sct.shot(output=output)
        print(f"Full screen screenshot saved as '{output}'")


def custom_area_screenshot(x, y, width, height, output="custom_area_screenshot.png"):
    """
        自定义区域截图
    """
    """Capture a screenshot of a custom area on the screen."""
    with mss.mss() as sct:
        box = {"top": y, "left": x, "width": width, "height": height}
        screenshot = sct.grab(box)
        
        mss.tools.to_png(screenshot.rgb, screenshot.size, output=output)
        print(f"Custom area screenshot saved as '{output}' for area {box}")
 

class TransparentScreenshotApp:
    """
        画出自定义区域
    """        
    def __init__(self, root):
        self.root = root
        self.root.title("Transparent Custom Area Screenshot")
        self.root.attributes('-fullscreen', True)
        self.root.attributes('-alpha', 0.3)  # Set initial transparency
        # self.root.bind("<Return>", self.on_entry)
        self.root.bind("<Motion>", self.move_press)
        self.root.bind("<ButtonPress-1>", self.on_button_press)
        self.root.bind("<B1-Motion>", self.on_move_press)
        self.root.bind("<ButtonRelease-1>", self.on_button_release)
        self.root.bind("<Return>", self.on_enter_press)

        self.start_x = None
        self.start_y = None
        self.end_x = None
        self.end_y = None
        self.rect = None
        
        self.x = None
        self.y = None
        self.width = None
        self.height = None
        
        self.screen_width = None
        self.screen_height = None
        self.canvas = None
        
        self.rect_id = None
        

    def move_press(self, event):
        """
            鼠标移动， 窗口设置透明度
        """
        self.start_x = self.root.winfo_pointerx() - self.root.winfo_rootx()
        self.start_y = self.root.winfo_pointery() - self.root.winfo_rooty()
        self.root.attributes('-alpha', 0.3)


    def on_button_press(self, event):
        """
            鼠标按下，初始化canvas，后续需要用它画出矩形
        """
        self.start_x = self.root.winfo_pointerx() - self.root.winfo_rootx()
        self.start_y = self.root.winfo_pointery() - self.root.winfo_rooty()
        # self.rect = self.root.create_rectangle(self.start_x, self.start_y, self.start_x, self.start_y, outline='red', width=2)
        self.root.attributes('-alpha', 0.4)  # Increase transparency when dragging
        self.canvas = tk.Canvas(root, bg='black')
        self.canvas.pack(fill=tk.BOTH, expand=True)
        self.screen_width = root.winfo_screenwidth()
        self.screen_height = root.winfo_screenheight()
        

    def on_move_press(self, event):
        """
            鼠标移动，画出矩形。移动鼠标时，当前一个矩形被画出，新的矩形覆盖旧的矩形
        """
        self.end_x = event.x
        self.end_y = event.y
        if self.rect_id is not None:
            self.canvas.delete(self.rect_id)
        self.rect_id = self.canvas.create_rectangle(self.start_x, self.start_y, self.end_x, self.end_y, 
                                                        outline='red', fill='black', width=4)
            
        # self.root.coords(self.rect, self.start_x, self.start_y, curX, curY)

    def on_button_release(self, event):
        """
            鼠标松开，画出矩形，弹出提示框
        """
        self.root.attributes('-alpha', 0.5)  # Reset transparency
        if self.start_x is not None and self.end_x is not None and self.start_y is not None and self.end_y is not None:
            tk.messagebox.showinfo("提示", "请选择按下回车键进行截图")
        # self.root.coords(self.rect, self.start_x, self.start_y, self.end_x, self.end_y)
        

    def on_enter_press(self, event):
        """
            按下回车键，获取矩形区域，进行截图
        """
        x = min(self.start_x, self.end_x)
        y = min(self.start_y, self.end_y)
        width = abs(self.end_x - self.start_x)
        height = abs(self.end_y - self.start_y)
        # print(f"Selected area: top={y}, left={x}, width={width}, height={height}")
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        root.attributes('-alpha',0.0) # Reset transparency 截图前界面恢复原始透明度
        self.root.destroy()
        print(f"Selected area: top={y}, left={x}, width={width}, height={height}")
        custom_x, custom_y = x, y
        custom_width, custom_height = width, height
        custom_area_screenshot(custom_x, custom_y, custom_width, custom_height)

if __name__ == "__main__":
    # full_screen_screenshot()
    
    root = tk.Tk()
    app = TransparentScreenshotApp(root)
    root.mainloop()
    
    # print(f"Selected area: top={app.y}, left={app.x}, width={app.width}, height={app.height}")
    # custom_x, custom_y = app.x, app.y
    # custom_width, custom_height = app.width, app.height
    # custom_area_screenshot(custom_x, custom_y, custom_width, custom_height)